# Ultralytics YOLO 🚀, AGPL-3.0 license

import argparse

import cv2
import numpy as np
import onnxruntime as ort
import torch
import yaml

from ultralytics.utils import ASSETS
from ultralytics.utils.checks import check_requirements, check_yaml


def yaml_load(file="data.yaml", append_filename=False):
    with open(file, errors="ignore", encoding="utf-8") as f:
        data = yaml.safe_load(f.read()) or {}
        if append_filename:
            data["yaml_file"] = str(file)
        return data


class YOLOv8:
    """YOLOv8目标检测模型类，用于处理推理和可视化。"""

    def __init__(self, onnx_model, input_image, confidence_thres, iou_thres):
        """
        初始化YOLOv8类的实例。

        参数:
            onnx_model: ONNX模型路径。
            input_image: 输入图像路径。
            confidence_thres: 过滤检测结果的置信度阈值。
            iou_thres: 非极大值抑制的IoU(交并比)阈值。
        """
        self.onnx_model = onnx_model
        self.input_image = input_image
        self.confidence_thres = confidence_thres
        self.iou_thres = iou_thres

        # 从COCO数据集加载类名
        self.classes = yaml_load(check_yaml("VOC-tt100k.yaml"))["names"]

        # 为类别生成颜色调色板
        self.color_palette = np.random.uniform(0, 255, size=(len(self.classes), 3))

    def draw_detections(self, img, box, score, class_id):
        """
        根据检测到的对象在输入图像上绘制边界框和标签。

        参数:
            img: 绘制检测结果的输入图像。
            box: 检测到的边界框。
            score: 对应的检测置信度。
            class_id: 检测对象的类别ID。

        返回:
            无
        """

        # 提取边界框的坐标
        x1, y1, w, h = box

        # 检索该类ID的颜色
        color = self.color_palette[class_id]

        # 在图像上绘制边界框
        cv2.rectangle(img, (int(x1), int(y1)), (int(x1 + w), int(y1 + h)), color, 2)

        # 创建带有类名和得分的标签文本
        label = f"{self.classes[class_id]}: {score:.2f}"

        # 计算标签文本的尺寸
        (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)

        # 计算标签文本的位置
        label_x = x1
        label_y = y1 - 10 if y1 - 10 > label_height else y1 + 10

        # 绘制填充矩形作为标签文本的背景
        cv2.rectangle(
            img, (label_x, label_y - label_height), (label_x + label_width, label_y + label_height), color, cv2.FILLED
        )

        # 在图像上绘制标签文本
        cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA)

    def preprocess(self):
        """
        在执行推理之前预处理输入图像。

        返回:
            image_data: 准备好进行推理的预处理图像数据。
        """
        # 使用OpenCV读取输入图像
        self.img = cv2.imread(self.input_image)

        # 获取输入图像的高度和宽度
        self.img_height, self.img_width = self.img.shape[:2]

        # 将图像颜色空间从BGR转换为RGB
        img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)

        # 调整图像大小以匹配输入形状
        img = cv2.resize(img, (self.input_width, self.input_height))

        # 通过除以255.0来归一化图像数据
        image_data = np.array(img) / 255.0

        # 转置图像以将通道维度作为第一维度
        image_data = np.transpose(image_data, (2, 0, 1))  # 通道优先

        # 扩展图像数据的维度以匹配预期的输入形状
        image_data = np.expand_dims(image_data, axis=0).astype(np.float32)

        # 返回预处理后的图像数据
        return image_data

    def postprocess(self, input_image, output):
        """
        对模型输出进行后处理，以提取边界框、得分和类别ID。

        参数:
            input_image (numpy.ndarray): 输入图像。
            output (numpy.ndarray): 模型的输出。

        返回:
            numpy.ndarray: 绘制了检测结果的输入图像。
        """

        # 转置并压缩输出以匹配预期形状
        outputs = np.transpose(np.squeeze(output[0]))

        # 获取输出数组中的行数
        rows = outputs.shape[0]

        # 用于存储检测结果的边界框、得分和类别ID的列表
        boxes = []
        scores = []
        class_ids = []

        # 计算边界框坐标的缩放因子
        x_factor = self.img_width / self.input_width
        y_factor = self.img_height / self.input_height

        # 遍历输出数组中的每一行
        for i in range(rows):
            # 从当前行提取类别得分
            classes_scores = outputs[i][4:]

            # 找出类别得分中的最大分数
            max_score = np.amax(classes_scores)

            # 如果最高分数超过置信度阈值
            if max_score >= self.confidence_thres:
                # 获取得分最高的类别ID
                class_id = np.argmax(classes_scores)

                # 从当前行提取边界框坐标
                x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]

                # 计算边界框的缩放后坐标
                left = int((x - w / 2) * x_factor)
                top = int((y - h / 2) * y_factor)
                width = int(w * x_factor)
                height = int(h * y_factor)

                # 将类别ID、得分和边界框坐标添加到相应的列表中
                class_ids.append(class_id)
                scores.append(max_score)
                boxes.append([left, top, width, height])

        # 应用非极大值抑制来过滤重叠的边界框
        indices = cv2.dnn.NMSBoxes(boxes, scores, self.confidence_thres, self.iou_thres)

        # 遍历非极大值抑制后选定的索引
        for i in indices:
            # 获取与索引对应的边界框、得分和类别ID
            box = boxes[i]
            score = scores[i]
            class_id = class_ids[i]

            # 在输入图像上绘制检测结果
            self.draw_detections(input_image, box, score, class_id)

        # 返回修改后的输入图像
        return input_image

    def main(self):
        """
        使用ONNX模型执行推理并返回带有绘制检测结果的输出图像。

        返回:
            output_img: 带有绘制检测结果的输出图像。
        """
        # 使用ONNX模型创建推理会话并指定执行提供程序
        session = ort.InferenceSession(self.onnx_model, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])

        # 获取模型输入
        model_inputs = session.get_inputs()

        # 存储输入形状以供后续使用
        input_shape = model_inputs[0].shape
        self.input_width = input_shape[2]
        self.input_height = input_shape[3]

        # 预处理图像数据
        img_data = self.preprocess()

        # 使用预处理后的图像数据运行推理
        outputs = session.run(None, {model_inputs[0].name: img_data})

        # 对输出进行后处理以获取输出图像。
        return self.postprocess(self.img, outputs)  # 输出图像


if __name__ == "__main__":
    # 创建参数解析器来处理命令行参数
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default=r"H:\vscode\work\ultralytics\ultralytics\weights\best.onnx", help="输入你的ONNX模型。")
    parser.add_argument("--img", type=str, default=r"H:\vscode\work\labelingPhoto\5.jpg", help="输入图像的路径。")
    parser.add_argument("--conf-thres", type=float, default=0.5, help="置信度阈值")
    parser.add_argument("--iou-thres", type=float, default=0.5, help="NMS IoU阈值")
    args = parser.parse_args()

    # 检查需求并选择适当的后端（CPU或GPU）
    check_requirements("onnxruntime-gpu" if torch.cuda.is_available() else "onnxruntime")

    # 使用指定的参数创建YOLOv8类的实例
    detection = YOLOv8(args.model, args.img, args.conf_thres, args.iou_thres)

    # 执行目标检测并获取输出图像
    output_image = detection.main()

    # 在窗口中显示输出图像
    cv2.namedWindow("输出", cv2.WINDOW_NORMAL)
    cv2.imshow("输出", output_image)

    # 等待按键退出
    cv2.waitKey(0)
